import torch.nn as nn
import torch.nn.functional as F
import torch

__all__ = ['FinerSeg', 'MSFinerSeg']


def initialize_weights(*models):
    for model in models:
        for m in model.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
                nn.init.kaiming_normal(m.weight)
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()


def conv3x3(in_channels, out_channels):
    return nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
                         nn.BatchNorm2d(out_channels),
                         nn.ReLU(inplace=True))


class SEModule(nn.Module):
    def __init__(self, channels, reduction=16):
        super(SEModule, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc1 = nn.Conv2d(channels, channels // reduction, kernel_size=1, padding=0)
        self.relu = nn.ReLU(inplace=True)
        self.fc2 = nn.Conv2d(channels // reduction, channels, kernel_size=1, padding=0)
        self.sigmoid = nn.Sigmoid()

    def forward(self, input):
        x = self.avg_pool(input)
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        x = self.sigmoid(x)
        return input * x


class Res2Block(nn.Module):
    """
    The Res2Net Block
    """

    def __init__(self, in_channels, out_channels, scales=4, stride=1, expansion=4):
        super(Res2Block, self).__init__()
        if out_channels % scales != 0:
            raise ValueError(
                "channels must be divisible by scales. out_channels:{0:d}, scales:{1:d}".format(out_channels, scales))
        self.sub_channels = out_channels // scales
        self.scales = scales
        self.expansion = expansion
        self.stride = stride
        self.conv1x1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, bias=False)
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=self.stride)
        self.bn1 = nn.BatchNorm2d(out_channels)
        if scales == 1:
            self.nums = 1
        else:
            self.nums = scales - 1
        convs = []
        bns = []
        for i in range(self.nums):
            convs.append(nn.Conv2d(self.sub_channels, self.sub_channels, kernel_size=3, stride=self.stride, padding=1, bias=False))
            bns.append(nn.BatchNorm2d(self.sub_channels))
        self.convs = nn.ModuleList(convs)
        self.bns = nn.ModuleList(bns)

        self.conv3 = nn.Conv2d(self.sub_channels * scales, self.sub_channels * self.expansion, kernel_size=1, bias=False)
        self.bn3 = nn.BatchNorm2d(out_channels * self.expansion)
        self.se = SEModule(self.sub_channels * self.expansion)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        residual = self.conv1x1(x)

        out = self.relu(self.bn1(self.conv1(x)))
        # We use torch.split operation to split the features into sub-block features in channels wise
        split_x = torch.split(out, self.sub_channels, dim=1)
        for i in range(self.nums):
            if i == 0:
                split = split_x[i]
            else:
                split = split + split_x[i]
            split = self.relu(self.bns[i](self.convs[i](split)))
            if i == 0:
                out = split
            else:
                out = torch.cat((out, split), dim=1)

        if self.scales > 1:
            out = torch.cat((out, split_x[self.nums]), dim=1)

        out = self.relu(self.conv3(out))
        out = self.se(out)
        out = out + residual
        out = self.relu(out)
        return out


class EncoderBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(EncoderBlock, self).__init__()
        self.conv = nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
                                  nn.BatchNorm2d(out_channels),
                                  nn.ReLU(inplace=True),
                                  nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
                                  nn.BatchNorm2d(out_channels),
                                  nn.ReLU(inplace=True))

    def forward(self, x):
        out = self.conv(x)
        return out


class UNet(nn.Module):
    """
    The UNet is used as a refine netowrk
    """

    def __init__(self, in_channels, out_channels):
        super(UNet, self).__init__()
        self.enc1 = EncoderBlock(in_channels, 32)
        self.enc2 = EncoderBlock(32, 32)
        self.enc3 = EncoderBlock(32, 32)
        self.enc4 = EncoderBlock(32, 32)
        self.bridge = EncoderBlock(32, 32)
        self.dec4 = EncoderBlock(64, 32)
        self.dec3 = EncoderBlock(64, 32)
        self.dec2 = EncoderBlock(64, 32)
        self.dec1 = EncoderBlock(64, 32)
        self.down = nn.MaxPool2d(kernel_size=2, stride=2)
        self.up = nn.Upsample(scale_factor=2, mode='bilinear')
        self.out = nn.Conv2d(32, out_channels, kernel_size=1, bias=False)
        initialize_weights(self)

    def forward(self, x):
        residual = x
        enc1 = self.enc1(x)
        down1 = self.down(enc1)
        enc2 = self.enc2(down1)
        down2 = self.down(enc2)
        enc3 = self.enc3(down2)
        down3 = self.down(enc3)
        enc4 = self.enc4(down3)
        down4 = self.down(enc4)
        bridge = self.bridge(down4)
        up4 = self.up(bridge)
        up4 = torch.cat((up4, enc4), 1)
        dec4 = self.dec4(up4)
        up3 = self.up(dec4)
        up3 = torch.cat((up3, enc3), 1)
        dec3 = self.dec3(up3)
        up2 = self.up(dec3)
        up2 = torch.cat((up2, enc2), 1)
        dec2 = self.dec2(up2)
        up1 = self.up(dec2)
        up1 = torch.cat((up1, enc1), 1)
        dec1 = self.dec1(up1)
        out = self.out(dec1)
        out = residual + out
        return out


class FinerSeg(nn.Module):
    def __init__(self, classes, channels, scales=4, expansion=4):
        """
        :param classes: output classes
        :param channels: input channels
        """
        super(FinerSeg, self).__init__()
        if scales != expansion:
            raise ValueError("scales and expansion must be the same")
        self.encoder1 = Res2Block(channels, 32, scales=scales, expansion=expansion)
        self.encoder2 = Res2Block(32, 64, scales=scales, expansion=expansion)
        self.encoder3 = Res2Block(64, 128, scales=scales, expansion=expansion)
        self.encoder4 = Res2Block(128, 256, scales=scales, expansion=expansion)
        self.bridge = Res2Block(256, 512, scales=scales, expansion=expansion)
        self.down = nn.MaxPool2d(kernel_size=2, stride=2)
        self.decoder4 = Res2Block(512, 256, scales=scales, expansion=expansion)
        self.decoder3 = Res2Block(256, 128, scales=scales, expansion=expansion)
        self.decoder2 = Res2Block(128, 64, scales=scales, expansion=expansion)
        self.decoder1 = Res2Block(64, 32, scales=scales, expansion=expansion)
        self.upsample4 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
        self.upsample3 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.upsample2 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
        self.upsample1 = nn.ConvTranspose2d(64, 32, kernel_size=2, stride=2)
        self.final = nn.Conv2d(32, classes, kernel_size=1, stride=1, bias=False)
        self.refine = UNet(classes, classes)
        initialize_weights(self)

    def forward(self, x):
        # Encoder Path
        enc1 = self.encoder1(x)
        down1 = self.down(enc1)
        enc2 = self.encoder2(down1)
        down2 = self.down(enc2)
        enc3 = self.encoder3(down2)
        down3 = self.down(enc3)
        enc4 = self.encoder4(down3)
        down4 = self.down(enc4)

        bridge = self.bridge(down4)

        # Decoder Path
        up4 = self.upsample4(bridge)
        up4 = torch.cat((up4, enc4), dim=1)
        dec4 = self.decoder4(up4)
        up3 = self.upsample3(dec4)
        up3 = torch.cat((up3, enc3), dim=1)
        dec3 = self.decoder3(up3)
        up2 = self.upsample2(dec3)
        up2 = torch.cat((up2, enc2), dim=1)
        dec2 = self.decoder2(up2)
        up1 = self.upsample1(dec2)
        up1 = torch.cat((up1, enc1), dim=1)
        dec1 = self.decoder1(up1)

        # Final output
        coarse_map = self.final(dec1)
        fine_map = self.refine(coarse_map)
        coarse_map = F.sigmoid(coarse_map)
        fine_map = F.sigmoid(fine_map)

        return coarse_map, fine_map
